import sys
import os
import torch
import torch.nn as nn
# 可视化张量图片
from utils import setup_seed
from tqdm import tqdm
from Loss.IIC_Loss import IIC_Loss
from torch.utils.tensorboard import SummaryWriter

class Trainer:
    def __init__(self, config, loader):
        self.writer = SummaryWriter(f"./project_datasave/{config['project_name']}/log")
        self.config = config
        self.train_cfg = config["trainer"]
        setup_seed(self.train_cfg["seed"])
        self.device = self.train_cfg["device"]
        self.train_time = 0
        self.train_loader = loader
        self.set_loss()
        self.set_net()
        self.load_weight()
        input_test_data = next(iter(self.train_loader))[0]
        self.writer.add_graph(self.model, input_test_data.to(self.device))

    def set_net(self):
        net_structure = self.train_cfg["net"]
        print("将使用%s网络结构进行训练" % net_structure)
        if (net_structure == 'cs3'):
            from Nets.cnn_single_head_3layer import CNNNet
            self.model = CNNNet(self.config).to(self.device)
        elif (net_structure == 'resnet18'):
            from Nets.Resnet import ResNet_18
            self.model = ResNet_18(self.config).to(self.device)
        elif (net_structure == 'cs4'):
            from Nets.cnn_single_head import CNNNet
            self.model = CNNNet(self.config).to(self.device)
        elif (net_structure == 'ms5'):
            from Nets.mlp_5layer import MLPNet
            self.model = MLPNet(self.config).to(self.device)
        elif (net_structure == 'cm3'):
            from Nets.cnn_multi_head import CNNNet
            self.model = CNNNet(self.config).to(self.device)
        elif (net_structure == 'mm5'):
            from Nets.mlp_multi_head_5layer import MLPNet
            self.model = MLPNet(self.config).to(self.device)
        elif (net_structure == 'mm7'):
            from Nets.mlp_multi_head_7layer import MLPNet
            self.model = MLPNet(self.config).to(self.device)
        elif (net_structure == 'cmp3'):
            from Nets.cnn_multi_head_pca import CNNNet
            self.model = CNNNet(self.config).to(self.device)
        elif (net_structure == 'csp3'):
            from Nets.cnn_single_head_pca import CNNNet
            self.model = CNNNet(self.config).to(self.device)
        if (self.train_cfg["opt"] == 'adam'):
            self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.train_cfg["lr"])

    def set_loss(self):
        loss_type = self.train_cfg["loss"]
        if (loss_type == "cross_entropy"):
            self.loss = nn.CrossEntropyLoss().to(self.device)
        elif (loss_type == "iic"):
            self.loss = IIC_Loss().to(self.device)

    def save_model(self, epoch):
        if not os.path.exists("./project_datasave/" + self.config["project_name"] + "/model/every_epochs"):
            os.makedirs("./project_datasave/" + self.config["project_name"] + "/model/every_epochs")
        # 保存权重模型
        torch.save(self.model.state_dict(), "./project_datasave/" + self.config["project_name"] + "/model/every_epochs/model_" + str(epoch) + ".pth")

    def load_weight(self):
        if (self.train_cfg["load_weight"] != "" and self.train_cfg["load_weight"] != None):
            model_path = self.train_cfg["load_weight"]
            print(f"将加载模型权重{model_path}......")
            self.model.load_state_dict(torch.load(model_path))

    def save_loss(self, loss_save_list):
        with open("./project_datasave/" + self.config["project_name"] + "/log/loss.txt", "w") as f:
            for loss_epoch in loss_save_list:
                f.write(str(loss_epoch) + "\n")
        # 绘制损失曲线
        import matplotlib.pyplot as plt
        plt.plot(loss_save_list)
        plt.xlabel("epoch")
        plt.ylabel("loss")
        plt.savefig("./project_datasave/" + self.config["project_name"] + "/figure/loss.png")

    def train():
        pass
